#include "fiber.h"

static bool debug = false;

namespace lzp {
// thread_local可以保证不同线程中的thread_local变量互不影响，即线程局部变量
// 当前线程正在运行的协程
static thread_local Fiber* t_fiber = nullptr;
// 当前线程的主协程
static thread_local std::shared_ptr<Fiber> t_thread_fiber = nullptr;
// 当前线程的调度协程，每个线程都独有一份，包括caller线程
static thread_local Fiber* t_scheduler_fiber = nullptr;


// 协程id
static std::atomic<uint64_t> s_fiber_id{0};
// 协程计数器
static std::atomic<uint64_t> s_fiber_count{0};

/**
 * @brief 设置当前正在运行的协程，即t_fiber
 */
void Fiber::SetThis(Fiber *f) {
    t_fiber = f;
}

/**
 * @brief 返回当前线程正在执行的协程
 * @details 如果当前线程还未创建协程，则创建线程的第一个协程
 * 且该协程为当前线程的主协程，其他协程都通过这个协程来调度
 * 也就是说，其他协程结束时，都要切回到主协程，由主协程重新选择新的协程进行resume
 * @attention 线程如果要创建协程，那么应该首先执行Fiber::GetThis()，以初始化主函数协程
 */
std::shared_ptr<Fiber> Fiber::GetThis() {
    if (t_fiber) {
        return t_fiber->shared_from_this();
    }

    // 创建主协程
    std::shared_ptr<Fiber> main_fiber(new Fiber());
    t_thread_fiber = main_fiber;
    t_scheduler_fiber = main_fiber.get(); // 除非主动设置 主协程默认为调度协程

    // assert检查条件是否为真，条件不为真则终止程序并输出一条错误信息
    assert(t_fiber == main_fiber.get());
    return t_fiber->shared_from_this();
}

/**
 * @brief 设置调度协程（默认为主协程）
 */
void Fiber::SetSchedulerFiber(Fiber* f) {
    t_scheduler_fiber = f;
}

/**
 * @brief 得到当前运行的协程id
 */
uint64_t Fiber::GetFiberId() {
    if (t_fiber) {
        return t_fiber->getId();
    }
    return (uint64_t) - 1;
}

/**
 * @brief 构造函数
 * @attention 无参构造函数只用于创建线程的第一个协程，即线程主函数对应的协程
 * 这个构造函数只会在GetThis()方法调用 
 */
Fiber::Fiber() {
    SetThis(this);
    m_state = RUNNING;

    if (getcontext(&m_ctx)) {
		std::cerr << "Fiber() failed\n";
		pthread_exit(NULL);
    }

    // 协程id从0开始
    m_id = s_fiber_id++;
    ++s_fiber_count;

    // for debug
    if(debug) std::cout << "Fiber(): main id = " << m_id << std::endl;
}

/** 
 * @brief 构造函数，用于创建用户协程
 * @param[in] cb 协程入口函数
 * @param[in] stacksize 栈大小
 * @param[in] run_in_scheduler 本协程是否参与调度器调度，默认为true 
 */
Fiber::Fiber(std::function<void()> cb, size_t stacksize, bool run_in_scheduler) : m_cb(cb), m_runInScheduler(run_in_scheduler) {
    m_state = READY;

    // 分配协程栈空间
    m_stacksize = stacksize ? stacksize : 128000;
    // 在堆上分配协程的栈空间
    m_stack = malloc(m_stacksize);
    
    if (getcontext(&m_ctx)) {
        std::cerr << "Fiber(std::function<void()> cb, size_t stacksize, bool run_in_scheduler) failed\n";
		pthread_exit(NULL);
    }

    m_ctx.uc_link = nullptr;
    m_ctx.uc_stack.ss_sp = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;
    makecontext(&m_ctx, &Fiber::MainFunc, 0);

    m_id = s_fiber_id++;
    ++s_fiber_count;

    // for debug
    if(debug) std::cout << "Fiber(): child id = " << m_id << std::endl;
}

/**
 * @brief 析构函数
 */
Fiber::~Fiber() {
    --s_fiber_count;
    if (m_stack) {
        free(m_stack);
    }
    if(debug) std::cout << "~Fiber(): id = " << m_id << std::endl;	
}

/**
 * @brief 重置协程状态和入口函数，复用栈空间，不重新创建栈
 * @param[in] cb 
 */
void Fiber::reset(std::function<void()> cb) {
    // 因为要复用栈空间，所以栈指针不能为空，并且协程状态要为终止态
    assert(m_stack != nullptr && m_state == TERM);

    m_state = READY;
    m_cb = cb;

    if (getcontext(&m_ctx)) {
        std::cerr << "reset() failed\n";
		pthread_exit(NULL);
    }

    m_ctx.uc_link = nullptr;
    m_ctx.uc_stack.ss_sp = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;
    makecontext(&m_ctx, &Fiber::MainFunc, 0);
}

/**
 * @brief 将当前协程切换到执行态
 * @details 当前协程和正在运行的协程进行交换，前者变为RUNNING，后者变为READY 
 */
void Fiber::resume() {
    // 要恢复的协程不能为终止态也不能为运行态
    assert(m_state != TERM && m_state != RUNNING);
    // 设置运行协程为当前协程
    SetThis(this);
    m_state = RUNNING;

    // 如果协程参与调度器调度，那么应该和调度器的主协程进行swap，而不是线程主协程
    if (m_runInScheduler) {
        if (swapcontext(&(t_scheduler_fiber->m_ctx), &m_ctx)) {
			std::cerr << "resume() to t_scheduler_fiber failed\n";
			pthread_exit(NULL);
        }
    }
    else {
        if (swapcontext(&(t_thread_fiber->m_ctx), &m_ctx)) {
			std::cerr << "resume() to t_thread_fiber failed\n";
			pthread_exit(NULL);
        }
    }
}

/**
 * @brief 当前任务让出执行权
 * @details 当前协程和上次resume时退到后台的协程进行交换，前者变为READY，后者变为RUNNING
 */
void Fiber::yield() {
    // 要让出cpu的协程只能为运行或者终止态，协程运行完后会自动yield一次，用于回到主协程，此时状态已为结束状态
    assert(m_state == RUNNING || m_state == TERM);

    // 没有终止的协程则重新设置为READY态
    if (m_state != TERM) {
        m_state = READY;
    }

    // 如果协程参与调度器调度，那么应该和调度器的主协程swap，而不是线程主协程
    if (m_runInScheduler) {
        SetThis(t_scheduler_fiber);
        if (swapcontext(&m_ctx, &(t_scheduler_fiber->m_ctx))) {
			std::cerr << "yield() to to t_scheduler_fiber failed\n";
			pthread_exit(NULL);
        }
    }
    else {
        SetThis(t_thread_fiber.get());
        if (swapcontext(&m_ctx, &(t_thread_fiber->m_ctx))) {
			std::cerr << "yield() to t_thread_fiber failed\n";
			pthread_exit(NULL);
        }
    }
}

/**
 * @brief 协程入口函数
 */
void Fiber::MainFunc() {
    // GetThis()的shared_from_this()让引用计数加1
    std::shared_ptr<Fiber> curr = GetThis();
    assert(curr != nullptr);

    // 执行协程入口函数
    curr->m_cb();
    curr->m_cb = nullptr;
    curr->m_state = TERM;

    // 运行完毕 -> 让出执行权
    auto raw_ptr = curr.get();
    // shared_ptr的reset()可以释放当前智能指针对象的所有权，即让引用计数减1，当引用计数为0时释放
    curr.reset();
    // 协程结束时自动yield，回到主协程
    raw_ptr->yield();
}

}
